/-
Copyright (c) 2025 Lean FRO, LLC. All rights reserved.
Released under Apache 2.0 license as described in the file LICENSE.
Authors: Paul Reichert
-/
module

prelude
public import Init.WF
import Init.Classical
import Init.Ext
import Init.NotationExtra

set_option doc.verso true
set_option linter.missingDocs true

/-!
This module provides a fixpoint combinator that combines advantages of `partial` and well-founded
recursion.

The combinator is similar to {lean}`WellFounded.fix`, but does
not require a well-foundedness proof. Therefore, it can be used in situations in which
a proof of termination is impossible or inconvenient. Given a well-foundedness proof, there is a
theorem guaranteeing that it is equal to {lean}`WellFounded.fix`.
-/

variable {α : Sort _} {β : α → Sort _} {γ : (a : α) → β a → Sort _}
  {C : α → Sort _} {C₂ : (a : α) → β a → Sort _} {C₃ : (a : α) → (b : β a) → γ a b → Sort _}

set_option doc.verso true

namespace WellFounded

/--
The function implemented as the loop {lean}`opaqueFix R F a = F a (fun a _ => opaqueFix R F a)`.
{lean}`opaqueFix R F` is the fixpoint of the functional {name}`F`, as long as the loop is
well-founded.

The loop might run forever depending on {name}`F`. It is opaque, i.e., it is impossible to prove
nontrivial properties about it.
-/
@[specialize]
partial def opaqueFix [∀ a, Nonempty (C a)] (R : α → α → Prop)
    (F : ∀ a, (∀ a', R a' a → C a') → C a) (a : α) : C a :=
  F a (fun a _ => opaqueFix R F a)

/-
SAFE assuming that the code generated by iteration over `F` is equivalent to the well-founded
fixpoint of `F` if `R` is well-founded.
-/
/--
A fixpoint combinator that can be used to construct recursive functions with an *extrinsic* proof
of termination.

Given a relation {name}`R` and a fixpoint functional {name}`F` which must be decreasing with respect
to {name}`R`, {lean}`extrinsicFix R F` is the recursive function obtained by having {name}`F` call
itself recursively.

If {name}`R` is not well-founded, {lean}`extrinsicFix R F` might run forever. In this case,
nothing interesting can be proved about the recursive function; it is opaque.

If {name}`R` _is_ well-founded, {lean}`extrinsicFix R F` is equivalent to
{lean}`WellFounded.fix _ F`, logically and regarding its termination behavior.
-/
@[implemented_by opaqueFix]
public def extrinsicFix [∀ a, Nonempty (C a)] (R : α → α → Prop)
    (F : ∀ a, (∀ a', R a' a → C a') → C a) (a : α) : C a :=
  open scoped Classical in
  if h : WellFounded R then
    h.fix F a
  else
    -- Return `opaqueFix R F a` so that `implemented_by opaqueFix` is sound.
    -- In effect, `extrinsicFix` is opaque if `WellFounded R` is false.
    opaqueFix R F a

public theorem extrinsicFix_eq_fix [∀ a, Nonempty (C a)] {R : α → α → Prop}
    {F : ∀ a, (∀ a', R a' a → C a') → C a}
    (wf : WellFounded R) {a : α} :
    extrinsicFix R F a = wf.fix F a := by
  simp only [extrinsicFix, dif_pos wf]

public theorem extrinsicFix_eq_apply [∀ a, Nonempty (C a)] {R : α → α → Prop}
    {F : ∀ a, (∀ a', R a' a → C a') → C a} (h : WellFounded R) {a : α} :
    extrinsicFix R F a = F a (fun a _ => extrinsicFix R F a) := by
  simp only [extrinsicFix, dif_pos h]
  rw [WellFounded.fix_eq]

/--
A fixpoint combinator that allows for deferred proofs of termination.

{lean}`extrinsicFix R F` is function implemented as the loop
{lean}`extrinsicFix R F a = F a (fun a _ => extrinsicFix R F a)`.

If the loop can be shown to be well-founded, {name}`extrinsicFix_eq_apply` proves that it satisfies the
fixpoint equation. Otherwise, {lean}`extrinsicFix R F` is opaque, i.e., it is impossible to prove
nontrivial properties about it.
-/
add_decl_doc extrinsicFix

/--
The function implemented as the loop
{lean}`opaqueFix₂ R F a b = F a b (fun a' b' _ => opaqueFix₂ R F a' b')`.
{lean}`opaqueFix₂ R F` is the fixpoint of the functional {name}`F`, as long as the loop is
well-founded.

The loop might run forever depending on {name}`F`. It is opaque, i.e., it is impossible to prove
nontrivial properties about it.
-/
@[specialize]
partial def opaqueFix₂ [∀ a b, Nonempty (C₂ a b)]
    (R : (a : α) ×' β a → (a : α) ×' β a → Prop)
    (F : (a : α) → (b : β a) → ((a' : α) → (b' : β a') → R ⟨a', b'⟩ ⟨a, b⟩ → C₂ a' b') → C₂ a b)
    (a : α) (b : β a) :
    C₂ a b :=
  F a b (fun a' b' _ => opaqueFix₂ R F a' b')

/-
SAFE assuming that the code generated by iteration over `F` is equivalent to the well-founded
fixpoint of `F` if `R` is well-founded.
-/
/--
A fixpoint combinator that can be used to construct recursive functions of arity two with an
*extrinsic* proof of termination.

Given a relation {name}`R` and a fixpoint functional {name}`F` which must be decreasing with respect
to {name}`R`, {lean}`extrinsicFix₂ R F` is the recursive function obtained by having {name}`F` call
itself recursively.

If {name}`R` is not well-founded, {lean}`extrinsicFix₂ R F` might run forever. In this case,
nothing interesting can be proved about the recursive function; it is opaque.

If {name}`R` _is_ well-founded, {lean}`extrinsicFix₂ R F` is equivalent to a well-founded recursive
function, logically and regarding its termination behavior.
-/
@[implemented_by opaqueFix₂]
public def extrinsicFix₂ [∀ a b, Nonempty (C₂ a b)]
    (R : (a : α) ×' β a → (a : α) ×' β a → Prop)
    (F : (a : α) → (b : β a) → ((a' : α) → (b' : β a') → R ⟨a', b'⟩ ⟨a, b⟩ → C₂ a' b') → C₂ a b)
    (a : α) (b : β a) :
    C₂ a b :=
  let F' (x : PSigma β) (G : (y : PSigma β) → R y x → C₂ y.1 y.2) : C₂ x.1 x.2 :=
    F x.1 x.2 (fun a b => G ⟨a, b⟩)
  extrinsicFix (C := fun x : PSigma β => C₂ x.1 x.2) R F' ⟨a, b⟩

public theorem extrinsicFix₂_eq_fix [∀ a b, Nonempty (C₂ a b)]
    {R : (a : α) ×' β a → (a : α) ×' β a → Prop}
    {F : ∀ a b, (∀ a' b', R ⟨a', b'⟩ ⟨a, b⟩ → C₂ a' b') → C₂ a b}
    (wf : WellFounded R) {a b} :
    extrinsicFix₂ R F a b = wf.fix (fun x G => F x.1 x.2 (fun a b h => G ⟨a, b⟩ h)) ⟨a, b⟩ := by
  rw [extrinsicFix₂, extrinsicFix_eq_fix wf]

public theorem extrinsicFix₂_eq_apply [∀ a b, Nonempty (C₂ a b)]
    {R : (a : α) ×' β a → (a : α) ×' β a → Prop}
    {F : (a : α) → (b : β a) → ((a' : α) → (b' : β a') → R ⟨a', b'⟩ ⟨a, b⟩ → C₂ a' b') → C₂ a b}
    (wf : WellFounded R) {a : α} {b : β a} :
    extrinsicFix₂ R F a b = F a b (fun a' b' _ => extrinsicFix₂ R F a' b') := by
  rw [extrinsicFix₂, extrinsicFix_eq_apply wf]
  rfl

/--
A fixpoint combinator that allows for deferred proofs of termination.

{lean}`extrinsicFix₂ R F` is function implemented as the loop
{lean}`extrinsicFix₂ R F a b = F a b (fun a' b' _ => extrinsicFix₂ R F a' b')`.

If the loop can be shown to be well-founded, {name}`extrinsicFix₂_eq_apply` proves that it satisfies the
fixpoint equation. Otherwise, {lean}`extrinsicFix₂ R F` is opaque, i.e., it is impossible to prove
nontrivial properties about it.
-/
add_decl_doc extrinsicFix₂

/--
The function implemented as the loop
{lean}`opaqueFix₃ R F a b c = F a b c (fun a' b' c' _ => opaqueFix₃ R F a' b' c')`.
{lean}`opaqueFix₃ R F` is the fixpoint of the functional {name}`F`, as long as the loop is
well-founded.

The loop might run forever depending on {name}`F`. It is opaque, i.e., it is impossible to prove
nontrivial properties about it.
-/
@[specialize]
public partial def opaqueFix₃ [∀ a b c, Nonempty (C₃ a b c)]
    (R : (a : α) ×' (b : β a) ×' γ a b → (a : α) ×' (b : β a) ×' γ a b → Prop)
    (F : ∀ (a b c), (∀ (a' b' c'), R ⟨a', b', c'⟩ ⟨a, b, c⟩ → C₃ a' b' c') → C₃ a b c)
    (a : α) (b : β a) (c : γ a b) :
    C₃ a b c :=
  F a b c (fun a b c _ => opaqueFix₃ R F a b c)

/-
SAFE assuming that the code generated by iteration over `F` is equivalent to the well-founded
fixpoint of `F` if `R` is well-founded.
-/
/--
A fixpoint combinator that can be used to construct recursive functions of arity three with an
*extrinsic* proof of termination.

Given a relation {name}`R` and a fixpoint functional {name}`F` which must be decreasing with respect
to {name}`R`, {lean}`extrinsicFix₃ R F` is the recursive function obtained by having {name}`F` call
itself recursively.

If {name}`R` is not well-founded, {lean}`extrinsicFix₃ R F` might run forever. In this case,
nothing interesting can be proved about the recursive function; it is opaque.

If {name}`R` _is_ well-founded, {lean}`extrinsicFix₃ R F` is equivalent to a well-founded recursive
function, logically and regarding its termination behavior.
-/
@[implemented_by opaqueFix₃]
public def extrinsicFix₃ [∀ a b c, Nonempty (C₃ a b c)]
    (R : (a : α) ×' (b : β a) ×' γ a b → (a : α) ×' (b : β a) ×' γ a b → Prop)
    (F : ∀ (a b c), (∀ (a' b' c'), R ⟨a', b', c'⟩ ⟨a, b, c⟩ → C₃ a' b' c') → C₃ a b c)
    (a : α) (b : β a) (c : γ a b) :
    C₃ a b c :=
  let F' (x : (a : α) ×' (b : β a) ×' γ a b) (G : (y : (a : α) ×' (b : β a) ×' γ a b) → R y x → C₃ y.1 y.2.1 y.2.2) :
      C₃ x.1 x.2.1 x.2.2 :=
    F x.1 x.2.1 x.2.2 (fun a b c => G ⟨a, b, c⟩)
  extrinsicFix (C := fun x : (a : α) ×' (b : β a) ×' γ a b => C₃ x.1 x.2.1 x.2.2) R F' ⟨a, b, c⟩

public theorem extrinsicFix₃_eq_fix [∀ a b c, Nonempty (C₃ a b c)]
    {R : (a : α) ×' (b : β a) ×' γ a b → (a : α) ×' (b : β a) ×' γ a b → Prop}
    {F : ∀ a b c, (∀ a' b' c', R ⟨a', b', c'⟩ ⟨a, b, c⟩ → C₃ a' b' c') → C₃ a b c}
    (wf : WellFounded R) {a b c} :
    extrinsicFix₃ R F a b c = wf.fix (fun x G => F x.1 x.2.1 x.2.2 (fun a b c h => G ⟨a, b, c⟩ h)) ⟨a, b, c⟩ := by
  rw [extrinsicFix₃, extrinsicFix_eq_fix wf]

public theorem extrinsicFix₃_eq_apply [∀ a b c, Nonempty (C₃ a b c)]
    {R : (a : α) ×' (b : β a) ×' γ a b → (a : α) ×' (b : β a) ×' γ a b → Prop}
    {F : ∀ (a b c), (∀ (a' b' c'), R ⟨a', b', c'⟩ ⟨a, b, c⟩ → C₃ a' b' c') → C₃ a b c}
    (wf : WellFounded R) {a : α} {b : β a} {c : γ a b} :
    extrinsicFix₃ R F a b c = F a b c (fun a b c _ => extrinsicFix₃ R F a b c) := by
  rw [extrinsicFix₃, extrinsicFix_eq_apply wf]
  rfl

/--
A fixpoint combinator that allows for deferred proofs of termination.

{lean}`extrinsicFix₃ R F` is function implemented as the loop
{lean}`extrinsicFix₃ R F a b c = F a b c (fun a b c _ => extrinsicFix₃ R F a b c)`.

If the loop can be shown to be well-founded, {name}`extrinsicFix₃_eq_apply` proves that it satisfies the
fixpoint equation. Otherwise, {lean}`extrinsicFix₃ R F` is opaque, i.e., it is impossible to prove
nontrivial properties about it.
-/
add_decl_doc extrinsicFix₃

end WellFounded
